import gradio as gr
import pickle
import numpy as np
import cv2

# Load the saved KNN model
with open('best_knn_model.pkl', 'rb') as f:
    knn_model = pickle.load(f)

def predict_digit(image):
    try:
        if isinstance(image, np.ndarray):
            img_array = image
        else:
            img_array = np.array(image)
        
        if len(img_array.shape) == 3:
            img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
        
        img_resized = cv2.resize(img_array, (8, 8))
        img_inverted = 255 - img_resized
        img_normalized = (img_inverted / 255.0) * 16.0
        img_flattened = img_normalized.flatten().reshape(1, -1)
        
        prediction = knn_model.predict(img_flattened)
        probability = np.max(knn_model.predict_proba(img_flattened))
        
        return f"Predicted digit: {prediction[0]} (Confidence: {probability:.2f})"
    except Exception as e:
        return f"Error: {str(e)}"

# Create Gradio interface
interface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Sketchpad(type="numpy", label="Draw a digit (0-9)"),
    outputs=gr.Textbox(label="Prediction Result"),
    title="Handwritten Digit Recognition with KNN",
    description="Draw a digit (0-9) in the sketchpad and see the KNN model prediction"
)

# Launch the interface
if __name__ == "__main__":
    print("Starting Gradio Web Application...")
    print("Creating public link...")
    interface.launch(share=True)
